# -*- coding: utf-8 -*-

import os
import pandas as pd
import oyaml as yaml


def get_c0_and_c1_video_names(root_folder):
    c0_folder = os.path.join(root_folder, "c0")
    c0_video_names = [name for name in os.listdir(c0_folder) if ".avi" in name]
    c1_folder = os.path.join(root_folder, "c1")
    c1_video_names = [name for name in os.listdir(c1_folder) if ".avi" in name]
    return c0_video_names, c1_video_names


def main():
    root_folder = r"F:\data\bus_videos\59092_00E0B453CE0B\2019_3_18"
    c0_video_names, c1_video_names = get_c0_and_c1_video_names(root_folder)

    door_excel_file_path = os.path.join(root_folder, "information", "label", "door_label.xlsx")
    data = pd.read_excel(door_excel_file_path)
    data.fillna(0, inplace=True)

    video_door_status = dict()
    for i, row in data.iterrows():
        video_name = row.iloc[0]+".avi"
        if video_name in c0_video_names:
            video_path = os.path.join("c0", video_name)
        else:
            video_path = os.path.join("c1", video_name)
        door_status = {"开门": list(), "关门": list()}
        last_status, last_frame = None, None
        for j in range(1, len(row)):
            if j == 1:
                pre_frame = 1
            else:
                pre_frame = int(row.iloc[j-1])
            frame = int(row.iloc[j])
            if j % 2 == 1:
                if pre_frame != 0 and frame != 0:
                    door_status["关门"].append((pre_frame, frame))
                    last_status = "开门"
                    last_frame = frame
            else:
                if pre_frame != 0 and frame != 0:
                    door_status["开门"].append((pre_frame, frame))
                    last_status = "关门"
                    last_frame = frame
        if last_status is None:
            door_status["关门"] = [0, 1000000]
        else:
            door_status[last_status].append((last_frame, 1000000))
        video_door_status[video_path] = door_status

    label_folder = os.path.join(root_folder, "information", "label")
    os.makedirs(label_folder, exist_ok=True)
    yaml_file_path = os.path.join(label_folder, "door.yaml")
    with open(yaml_file_path, "w", encoding="utf-8") as yaml_file:
        yaml.dump(
            {"root": root_folder, "video_door_status": video_door_status},
            yaml_file,
            allow_unicode=True,
            default_flow_style=False)

if __name__ == "__main__":
    main()
